/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "chatglm6b_add_operation.h"
#include <iostream>
#include <sstream>
#include <unistd.h>
#include <cstring>
#include <syscall.h>
#include <securec.h>

#ifndef __CCE_KT_TEST__
#include "acl/acl.h"
extern void AddCustomDo(uint32_t coreDim, void* l2ctrl, void* stream, uint8_t* x, uint8_t* y, uint8_t* z,
    uint8_t* workspace, uint8_t* tiling);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void AddCustom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling);
#endif

namespace atb_speed {
static uint64_t g_operationId = 0;
const size_t HOST_TILING_BUFFER_DEFAULT_SIZE = 10240;
constexpr uint32_t MAX_PROFILING_FUNC_NAME = 2;
static constexpr int32_t DIM_1 = 1;
static constexpr int32_t DIM_2 = 2;
static constexpr int32_t DIM_3 = 3;
static constexpr int32_t DIM_4 = 4;

ChatGlm6BAddOperation::ChatGlm6BAddOperation(const std::string &name) : name_(name) {}

ChatGlm6BAddOperation::~ChatGlm6BAddOperation() {}

std::string ChatGlm6BAddOperation::GetName() const
{
    return name_;
}

atb::Status ChatGlm6BAddOperation::InferShape(const atb::SVector<atb::TensorDesc> &inTensorDescs,
    atb::SVector<atb::TensorDesc> &outTensorDescs) const
{
    return 0;
}

uint32_t ChatGlm6BAddOperation::GetInputNum() const
{
    return DIM_2;
}

uint32_t ChatGlm6BAddOperation::GetOutputNum() const
{
    return DIM_1;
}

atb::Status ChatGlm6BAddOperation::Setup(const atb::VariantPack &variantPack, uint64_t &workspaceSize, atb::Context *context)
{
    const size_t usrWorkspaceSize = 4096;
    const size_t sysWorkspaceSize = 16 * 1024 * 1024;

    workspaceSize = sysWorkspaceSize + usrWorkspaceSize;

    return 0;
}

atb::Status ChatGlm6BAddOperation::Execute(const atb::VariantPack &variantPack, uint8_t *workspace,
    uint64_t workspaceSize, atb::Context* context)
{
    const uint32_t blockDim = 8;
    const uint32_t oneRepeatCalcount = 128;  // fixed
    const uint32_t blockDimImm = 8;
    const uint32_t tileNumImm = 8;
    const uint32_t doubleBufferImm = 2;  // fixed
    const uint32_t totalLengthImm = blockDimImm * oneRepeatCalcount * tileNumImm * doubleBufferImm;
    const size_t outputByteSize = blockDim * 2048 * sizeof(uint16_t);
    const size_t tilingSize = 2 * sizeof(uint32_t);
    uint8_t *tilingHost = nullptr;
    uint8_t *tilingDevice = nullptr;
    CheckAcl(aclrtMallocHost((void**)(&tilingHost), tilingSize));
    
    if (memcpy_s(tilingHost, tilingSize, &totalLengthImm, sizeof(uint32_t)) != 0) {
        return 1;
    }
    if (memcpy_s(tilingHost + sizeof(uint32_t), tilingSize - sizeof(uint32_t), &tileNumImm, sizeof(uint32_t)) != 0) {
        return 1;
    }
    
    CheckAcl(aclrtMalloc((void**)&tilingDevice, tilingSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CheckAcl(aclrtMemcpy(tilingDevice, tilingSize, tilingHost, tilingSize, ACL_MEMCPY_HOST_TO_DEVICE));

    AddCustomDo(blockDim, nullptr, context->GetExecuteStream(),
        static_cast<uint8_t*>(variantPack.inTensors.at(0).deviceData),
        static_cast<uint8_t*>(variantPack.inTensors.at(1).deviceData),
        static_cast<uint8_t*>(variantPack.outTensors.at(0).deviceData), workspace, tilingDevice);
    
    CheckAcl(aclrtSynchronizeStream(context->GetExecuteStream()));

    CheckAcl(aclrtFree(tilingDevice));
    CheckAcl(aclrtFreeHost(tilingHost));

    return 0;
}

aclError ChatGlm6BAddOperation::CheckAcl(aclError ret)
{
    if (ret != ACL_ERROR_NONE) {
        std::cerr << __FILE__ << ":" << __LINE__ << " aclError:" << ret << std::endl;
    }
    return ret;
}
} // namespace atb_speed
